"""
By Jiyuan Liu (liujiyuan13@163.com), Jun. 28, 2021.
All rights reserved.
"""

import torch
from torch import nn
import torch.nn.init as init
import numpy as np
import mat73
import time
from util import *
import sklearn.preprocessing as prep

if __name__ == "__main__":

    # device config
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # dir config
    # data_name = 'Android_Mischief_V2'                    # no use
    # data_name = 'CTU_Encrypted_Malware_Traffic'          # no use
    # data_name = 'CTU_Encrypted_Malware_Traffic_2_class'  # no use
    # data_name = 'CTU_Normal_vs_Adw'
    # data_name = 'CTU_Normal_vs_Drp'
    # data_name = 'CTU_Normal_vs_Rtk'
    # data_name = 'CTU_Normal_vs_Susp'
    # data_name = 'CTU_Normal_vs_Trj'                      # no use
    data_name = 'CTU_Normal_vs_Adw_Drp_Rtk_Susp'
    print('# ', data_name)
    # data_dir = 'D:/Work/datasets/mData/maldata/amd/'     # no use
    data_dir = 'data/'
    model_dir = 'model/'
    split_dir = 'split/'
    res_dir = 'res/'

    # load data
    print('- load data ...')
    X, Y = load_data(data_dir + data_name + '.mat')
    # concatenate
    X = np.hstack(X)
    fea_dim, num_sample, num_class = X.shape[1], Y.shape[0], np.unique(Y).shape[0]

    # train-test split
    print('- split data ...')
    load_split = True
    if load_split:
        ind = np.load(split_dir + data_name + '_split.npz')
        ind_tr, ind_te = ind['ind_tr'], ind['ind_te']
    else:
        ind_tr, ind_te = gen_split(Y, tr_ratio=0.8)
        np.savez(split_dir + data_name + '_split.npz', ind_tr=ind_tr, ind_te=ind_te)
    X_tr, gt_tr = X[ind_tr,:], Y[ind_tr]
    X_te, gt_te = X[ind_te,:], Y[ind_te]


    # SVM
    print('- SVM')
    start = time.process_time()
    from sklearn import svm
    clf = svm.SVC(decision_function_shape='ovr')
    clf.fit(X_tr, gt_tr)
    y_te = clf.predict(X_te)
    fpr, fnr, error_rate = metric(gt_te, y_te, labels=[1, 0], percent=True)
    end = time.process_time()
    print('-- ts: {:.3f}s, fpr: {:.3f}, fnr: {:.3f}, error_rate: {:.3f}'.format(end - start, fpr, fnr, error_rate))
    # save result
    np.savez(res_dir + data_name + '_svm_res.npz', data_name=data_name, gt_te=gt_te, y_te=y_te,
             fnr=fnr, fpr=fpr, error_rate=error_rate, ts=end-start)

    # kNN
    print('- kNN')
    start = time.process_time()
    from sklearn import neighbors
    clf = neighbors.KNeighborsClassifier(num_class)
    clf.fit(X_tr, gt_tr)
    y_te = clf.predict(X_te)
    fpr, fnr, error_rate = metric(gt_te, y_te, labels=[1, 0], percent=True)
    end = time.process_time()
    print('-- ts: {:.3f}s, fpr: {:.3f}, fnr: {:.3f}, error_rate: {:.3f}'.format(end - start, fpr, fnr, error_rate))
    # save result
    np.savez(res_dir + data_name + '_knn_res.npz', data_name=data_name, gt_te=gt_te, y_te=y_te,
             fnr=fnr, fpr=fpr, error_rate=error_rate, ts=end - start)

    # # Guassian Process
    # print('- Guassian Process')
    # start = time.process_time()
    # from sklearn import gaussian_process
    # clf = gaussian_process.GaussianProcessClassifier()
    # clf.fit(X_tr, gt_tr)
    # y_te = clf.predict(X_te)
    # fpr, fnr, error_rate = metric(gt_te, y_te, labels=[1, 0], percent=True)
    # end = time.process_time()
    # print('-- ts: {:.3f}s, fpr: {:.3f}, fnr: {:.3f}, error_rate: {:.3f}'.format(end - start, fpr, fnr, error_rate))
    # # save result
    # np.savez(res_dir + data_name + '_gp_res.npz', data_name=data_name, gt_te=gt_te, y_te=y_te,
    #          fnr=fnr, fpr=fpr, error_rate=error_rate, ts=end - start)

    # Decision Tree
    print('- Decision Tree')
    start = time.process_time()
    from sklearn import tree
    clf = tree.DecisionTreeClassifier()
    clf.fit(X_tr, gt_tr)
    y_te = clf.predict(X_te)
    fpr, fnr, error_rate = metric(gt_te, y_te, labels=[1, 0], percent=True)
    end = time.process_time()
    print('-- ts: {:.3f}s, fpr: {:.3f}, fnr: {:.3f}, error_rate: {:.3f}'.format(end - start, fpr, fnr, error_rate))
    # save result
    np.savez(res_dir + data_name + '_dt_res.npz', data_name=data_name, gt_te=gt_te, y_te=y_te,
             fnr=fnr, fpr=fpr, error_rate=error_rate, ts=end - start)
